package sjc;

import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.leader.LeaderSelector;
import org.apache.curator.framework.recipes.leader.LeaderSelectorListenerAdapter;
import org.apache.curator.retry.ExponentialBackoffRetry;

import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class LeaderElectionTest {
    private static final int SESSION_TIMEOUT = 30 * 1000;      //会话超时时间
    private static final int CONNECTION_TIMEOUT = 3 * 1000;    //连接超时时间
    private static final int CLIENT_NUMBER = 10;               //客户端数量

    //ZooKeeper服务地址
    private static final String SERVER = "192.168.72.133:2181,192.168.72.134:2181,192.168.72.135:2181";
    private static final String PATH = "/leader_election";

    public static void main(String[] args) throws IOException, InterruptedException {
        //创建连接实例
        RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
        //自定义LeaderSelectorListenerAdapter实例集合
        List<CustomLeaderSelectorListenerAdapter> leaderSelectorListenerList = new ArrayList<CustomLeaderSelectorListenerAdapter>();

        //创建 CuratorFrameworkImpl实例
        CuratorFramework client = CuratorFrameworkFactory.newClient(SERVER, SESSION_TIMEOUT, CONNECTION_TIMEOUT, retryPolicy);
        client.start();

        for(int i=0;i<CLIENT_NUMBER;i++){
            //创建LeaderSelectorListenerAdapter实例
            CustomLeaderSelectorListenerAdapter leaderSelectorListener =  new CustomLeaderSelectorListenerAdapter(client, PATH, "Client #" + i);

            leaderSelectorListener.start();
            leaderSelectorListenerList.add(leaderSelectorListener);
        }

        //暂停当前线程，防止单元测试结束，可以让leader选举过程持续进行
        TimeUnit.SECONDS.sleep(600);

    }
}

class CustomLeaderSelectorListenerAdapter  extends LeaderSelectorListenerAdapter implements Closeable {

    private String name;
    private LeaderSelector leaderSelector;
    public AtomicInteger leaderCount = new AtomicInteger();

    public CustomLeaderSelectorListenerAdapter(CuratorFramework client, String path, String name) {
        this.name = name;
        this.leaderSelector = new LeaderSelector(client, path, this);
        // 自动重新排队, 该方法的调用可以确保此实例在释放领导权后还可能获得领导权
        leaderSelector.autoRequeue();
    }

    public void start() throws IOException {
        leaderSelector.start();
    }

    @Override
    public void close() throws IOException {
        leaderSelector.close();
    }

    /**
     * 获取领导权
     * @param client the client
     * @throws Exception
     */
    @Override
    public void takeLeadership(CuratorFramework client) throws Exception {
        final int waitSeconds = 5;
        System.out.println(name + "成为当前leader");
        System.out.println(name + " 之前成为leader的次数：" + leaderCount.getAndIncrement() + "次");

        //TODO 其他业务代码

        try{
            //等待5秒后放弃领导权（模拟业务执行过程）
            Thread.sleep(TimeUnit.SECONDS.toMillis(waitSeconds));
        }catch ( InterruptedException e ){
            System.err.println(name + "已被中断");
            Thread.currentThread().interrupt();
        }finally{
            System.out.println(name + "放弃领导权\n");
        }
    }
}
